/*
 * Copyright 2024-2025 Embabel Software, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.embabel.agent.a2a.server.support

import com.fasterxml.jackson.databind.ObjectMapper
import io.a2a.spec.Message
import io.a2a.spec.SendStreamingMessageResponse
import io.a2a.spec.StreamingEventKind
import io.a2a.spec.Task
import io.a2a.spec.TaskArtifactUpdateEvent
import io.a2a.spec.TaskStatusUpdateEvent
import org.slf4j.LoggerFactory
import org.springframework.http.MediaType
import org.springframework.stereotype.Service
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.Executors
import java.util.concurrent.ScheduledExecutorService
import java.util.concurrent.TimeUnit

/**
 * Handles streaming functionality for A2A messages
 */
@Service
class A2AStreamingHandler(
    private val objectMapper: ObjectMapper,
    private val taskStateManager: TaskStateManager
) {
    private val logger = LoggerFactory.getLogger(A2AStreamingHandler::class.java)
    private val activeStreams = ConcurrentHashMap<String, SseEmitter>()
    private val scheduler: ScheduledExecutorService = Executors.newScheduledThreadPool(1)

    /**
     * Creates a new SSE stream for the given stream ID
     * @param taskId Optional task ID to associate with this stream
     */
    fun createStream(streamId: String, taskId: String? = null, contextId: String? = null): SseEmitter {
        logger.info("Creating SSE stream for streamId: {}, taskId: {}", streamId, taskId)

        val emitter = SseEmitter(Long.MAX_VALUE)
        activeStreams[streamId] = emitter

        // Register task if taskId is provided
        if (taskId != null && contextId != null) {
            taskStateManager.registerTask(taskId, contextId, streamId)
        }

        emitter.onCompletion {
            logger.info("Stream completed for streamId: {}", streamId)
            activeStreams.remove(streamId)
        }

        emitter.onTimeout {
            logger.info("Stream timed out for streamId: {}", streamId)
            activeStreams.remove(streamId)
        }

        return emitter
    }

    /**
     * Resubscribes to an existing task by creating a new stream and replaying events
     */
    fun resubscribeToTask(taskId: String, newStreamId: String): SseEmitter {
        logger.info("Resubscribing to task: taskId={}, newStreamId={}", taskId, newStreamId)

        // Check if task exists
        if (!taskStateManager.taskExists(taskId)) {
            throw IllegalArgumentException("Task not found: $taskId")
        }

        val taskInfo = taskStateManager.getTaskInfo(taskId)!!
        val emitter = createStream(newStreamId, taskId, taskInfo.contextId)

        // Update the stream ID for this task
        taskStateManager.updateStreamId(taskId, newStreamId)

        // Replay all events for this task
        Thread.startVirtualThread {
            try {
                taskInfo.events.forEach { event ->
                    sendStreamEvent(newStreamId, event, taskId)
                }

                // If task is completed, close the stream
                if (!taskStateManager.isTaskActive(taskId)) {
                    closeStream(newStreamId)
                }
            } catch (e: Exception) {
                logger.error("Error replaying events for task {}", taskId, e)
                emitter.completeWithError(e)
            }
        }

        return emitter
    }

    /**
     * Sends a streaming event to the specified stream and records it for the task
     */
    fun sendStreamEvent(streamId: String, event: StreamingEventKind, taskId: String? = null) {
        val emitter = activeStreams[streamId] ?: run {
            logger.warn("No active stream found for streamId: {}", streamId)
            return
        }

        try {
            // Record event in task state manager if taskId is provided
            taskId?.let { tid ->
                taskStateManager.recordEvent(tid, event)
            }

            // All events must be wrapped in SendStreamingMessageResponse per A2A protocol
            val response = SendStreamingMessageResponse(
                "2.0",
                streamId,
                event,
                null
            )

            val eventData = SseEmitter.event()
                .data(objectMapper.writeValueAsString(response), MediaType.APPLICATION_JSON)
            emitter.send(eventData)
        } catch (e: Exception) {
            logger.error("Error sending stream event", e)
            emitter.completeWithError(e)
        }
    }

    /**
     * Closes the specified stream
     */
    fun closeStream(streamId: String) {
        val emitter = activeStreams.remove(streamId)
        emitter?.complete()
    }

    /**
     * Shuts down the streaming handler
     */
    fun shutdown() {
        scheduler.shutdown()
        try {
            if (!scheduler.awaitTermination(5, TimeUnit.SECONDS)) {
                scheduler.shutdownNow()
            }
        } catch (e: InterruptedException) {
            scheduler.shutdownNow()
        }
    }
}
